{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Matplotlib: colormap transformations\n",
    "======================================================================\n",
    "\n",
    "Operating on color vectors\n",
    "--------------------------\n",
    "\n",
    "Ever wanted to manipulate an existing colormap? Here is a routine to apply a function to the look up table of a colormap:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def cmap_map(function, cmap):\n",
    "    \"\"\" Applies function (which should operate on vectors of shape 3: [r, g, b]), on colormap cmap.\n",
    "    This routine will break any discontinuous points in a colormap.\n",
    "    \"\"\"\n",
    "    cdict = cmap._segmentdata\n",
    "    step_dict = {}\n",
    "    # Firt get the list of points where the segments start or end\n",
    "    for key in ('red', 'green', 'blue'):\n",
    "        step_dict[key] = list(map(lambda x: x[0], cdict[key]))\n",
    "    step_list = sum(step_dict.values(), [])\n",
    "    step_list = np.array(list(set(step_list)))\n",
    "    # Then compute the LUT, and apply the function to the LUT\n",
    "    reduced_cmap = lambda step : np.array(cmap(step)[0:3])\n",
    "    old_LUT = np.array(list(map(reduced_cmap, step_list)))\n",
    "    new_LUT = np.array(list(map(function, old_LUT)))\n",
    "    # Now try to make a minimal segment definition of the new LUT\n",
    "    cdict = {}\n",
    "    for i, key in enumerate(['red','green','blue']):\n",
    "        this_cdict = {}\n",
    "        for j, step in enumerate(step_list):\n",
    "            if step in step_dict[key]:\n",
    "                this_cdict[step] = new_LUT[j, i]\n",
    "            elif new_LUT[j,i] != old_LUT[j, i]:\n",
    "                this_cdict[step] = new_LUT[j, i]\n",
    "        colorvector = list(map(lambda x: x + (x[1], ), this_cdict.items()))\n",
    "        colorvector.sort()\n",
    "        cdict[key] = colorvector\n",
    "\n",
    "    return matplotlib.colors.LinearSegmentedColormap('colormap',cdict,1024)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets try it out: I want a jet colormap, but lighter, so that I can plot\n",
    "things on top of it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "light_jet = cmap_map(lambda x: x/2 + 0.5, matplotlib.cm.jet)\n",
    "\n",
    "x, y = np.mgrid[1:2, 1:10:0.01]\n",
    "plt.figure(figsize=[15, 1])\n",
    "plt.imshow(y, cmap=light_jet, aspect='auto')\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](files/attachments/Matplotlib_ColormapTransformations/jet_light.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly, if a darker jet colormap is desired:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "dark_jet = cmap_map(lambda x: x*0.75, matplotlib.cm.jet)\n",
    "\n",
    "x, y = np.mgrid[1:2, 1:10:0.01]\n",
    "plt.figure(figsize=[15, 1])\n",
    "plt.imshow(y, cmap=dark_jet, aspect='auto')\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](files/attachments/Matplotlib_ColormapTransformations/jet_dark.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a comparison, this is what the original jet looks like:\n",
    "![](files/attachments/Matplotlib_ColormapTransformations/jet_original.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Operating on indices\n",
    "--------------------\n",
    "\n",
    "OK, but what if you want to change the indices of a colormap, but not\n",
    "its colors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def cmap_xmap(function, cmap):\n",
    "    \"\"\" Applies function, on the indices of colormap cmap. Beware, function\n",
    "    should map the [0, 1] segment to itself, or you are in for surprises.\n",
    "\n",
    "    See also cmap_xmap.\n",
    "    \"\"\"\n",
    "    cdict = cmap._segmentdata\n",
    "    function_to_map = lambda x : (function(x[0]), x[1], x[2])\n",
    "    for key in ('red','green','blue'):\n",
    "        cdict[key] = map(function_to_map, cdict[key])\n",
    "        cdict[key].sort()\n",
    "        assert (cdict[key][0]<0 or cdict[key][-1]>1), \"Resulting indices extend out of the [0, 1] segment.\"\n",
    "\n",
    "    return matplotlib.colors.LinearSegmentedColormap('colormap',cdict,1024)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Discrete colormap\n",
    "-----------------\n",
    "\n",
    "Here is how you can discretize a continuous colormap."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def cmap_discretize(cmap, N):\n",
    "    \"\"\"Return a discrete colormap from the continuous colormap cmap.\n",
    "    \n",
    "        cmap: colormap instance, eg. cm.jet. \n",
    "        N: number of colors.\n",
    "    \"\"\"\n",
    "    if type(cmap) == str:\n",
    "        cmap = get_cmap(cmap)\n",
    "    colors_i = np.concatenate((np.linspace(0, 1., N), (0.,0.,0.,0.)))\n",
    "    colors_rgba = cmap(colors_i)\n",
    "    indices = np.linspace(0, 1., N+1)\n",
    "    cdict = {}\n",
    "    for ki, key in enumerate(('red','green','blue')):\n",
    "        cdict[key] = [(indices[i], colors_rgba[i-1,ki], colors_rgba[i,ki]) for i in range(N+1)]\n",
    "    # Return colormap object.\n",
    "    return matplotlib.colors.LinearSegmentedColormap(cmap.name + \"_%d\"%N, cdict, 1024)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So for instance, consider a discretized jet colormap with 6 colors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "discretized_jet = cmap_discretize(matplotlib.cm.jet, 6)\n",
    "\n",
    "x, y = np.mgrid[1:2, 1:10:0.01]\n",
    "plt.figure(figsize=[15, 1])\n",
    "plt.imshow(y, cmap=discretized_jet, aspect='auto')\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](files/attachments/Matplotlib_ColormapTransformations/jet_discretized.png)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
